#ifndef RSIM_STATE_H
#define RSIM_STATE_H

/* -------------------------------------------------------------------------- *
 * File: State.h                                                              *
 * Authors: Kishor Bhalerao                                                   *
 * Email : kishor8dm@gmail.com                                                *
 * Contributors:                                                              *
 * Email:                                                                     *
 * copyright (c) 2010 Authors.                                                *
 *                                                                            *
 * This program is free software: you can redistribute it and/or modify it    *
 * under the terms of the GNU General Public License as published by the Free *
 * Software Foundation, either version 3 of the License, or any later version.*
 *                                                                            *
 * This program is distributed in the hope that it will be useful, but        *
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY *
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License   *
 * for more details. (http://www.gnu.org/licenses/)                           *
 * -------------------------------------------------------------------------- */

#include "List.h"
#include "Index.h"
#include "Vector.h"
#include "LinAlgDecl.h"
#include "PrescribedMotion.h"
#include "Utils.h"
#include "Time.h"

#include <assert.h>


namespace RSIM{

class BodyIndex;
class MultibodySystem;
class PrescribedMotion;
class SystemData;
/** This class is introduced to allow easy access to the state vector of the system.
State stores the entire state vector in a contiguous block of memory. This is distinct
from the Q's stored in each body. That is changing Q's i
Q's of each joint. */
        
class State{
        private:
                // This was only added as an afterthought. Perhaps state needs to be updated
                // based on what is happening elsewhere.
                SystemData *SysData_;
                
                TIME *time_;
                
                /** Stores the q's for each body */
                double *Q_;		
                
                /** Stores the qdot's for each body */
                double *QDot_;
                
                /** Stores the u's for each body */
                double *U_;
                
                /** Stores the udots's for each body */
                double *UDot_;
                
                /** Total number of Q's associated with multibody system */ 
                int nQ_;
                
                /** Total U's associated with the multibody system */
                int nU_;
                
                /** Stores number of Q's associated with each body */
                int *BodyNQ_;
                
                /** Stores number of U's associated with each body */
                int *BodyNU_;
                
                /** Total nodes in the system*/
                const int nNodes_;				
                
                List <int> QIndexForNode_;
                
                List <int> UIndexForNode_;
                                
                List <PrescribedMotion *> PM_;
                
                List <int> PrescribedMotionIndex_;
                
                /** Stores the indices of quaternions in use.*/
                List <int> QuaternionIndex_;
                
                /** This sets up memory and initializes Q_, U_ from the multibody system.*/
                void initialize(MultibodySystem& system);			
                                
                // These are illegal functions right now.
                /** To be implemented */
                void operator=(const State&);				
                
                /** To be implemented */
                State();
                
                /** To be implemented */
                State(const State& St);
                
        public:
                State(MultibodySystem& system);

                /** Clear allocated Memory */
                ~State();
                
                /** Return the location of Q corresponding to Node specified by NodeIndex from Q Cache */
                double* wgetQForNode(const int& NodeIndex);
                
                double* wgetQDotForNode(const int& NodeIndex);		
                
                /** Return the location of U corresponding to Node specified by NodeIndex from U Cache */
                double* wgetUForNode(const int& NodeIndex);
                
                /** Return the location of UDot corresponding to Node specified by NodeIndex from UDot Cache.
                Note that du/dt = udot. UIndexForNode_ can be used directly. For joints employing quaternions
                dq/dt != u */
                double* wgetUDotForNode(const int& NodeIndex);
                
                /** During integration, the norm of quaternions will certainly drift from unity. This function normalizes
                the quaternions in use. The locations of Quaternions are stored in QuaternionIndex_ */
                void normalizeQuaternions();
                
                /** Return the location of Q corresponding to Node specified by NodeIndex from Q Cache */
                const double* getQForNode(const int& NodeIndex)const;		
                
                /** returns NULL if Quaternions are not in use. Note. QDot is simply the time derivative 
                of hte quaternions .Even for a free joint, the length of QDot is 4. */
                const double* getQDotForNode(const int& NodeIndex)const;
                
                /** Return the location of U corresponding to Node specified by NodeIndex from U Cache */
                const double* getUForNode(const int& NodeIndex)const;
                
                /** Return the location of UDot corresponding to Node specified by NodeIndex from UDot Cache.
                Note that du/dt = udot. UIndexForNode_ can be used directly. For joints employing quaternions
                dq/dt != u */
                const double* getUDotForNode(const int& NodeIndex)const;
                
                Vector getSystemQAsVector()const;
                
                /** Return Q of the system as a vector. Debug function. Excess copy constructor on each call.*/
                Vector getQAsVectorForNode(const int& NodeIndex)const;
                
                /** Return U of the system as a vector. Debug function. Excess copy constructor on each call.*/
                const Vector getUAsVectorForNode(const int& NodeIndex)const;
                
                /** Return QDot of the system as a vector. Debug function. Excess copy constructor on each call.*/
                const Vector getQDotAsVectorForNode(const int& NodeIndex)const;
                
                /** Return UDot of the system as a vector. Debug function. Excess copy constructor on each call.*/
                const Vector getUDotAsVectorForNode(const int& NodeIndex)const;
                
                /** Total number of Q's associated with multibody system */ 
                const int getNQ() const;
                
                /** Total number of U's associated with multibody system */ 
                const int getNU() const;
                        
                const double* getQDot()const;
                
                const double* getUDot()const;
                
                void setY(const double *y_ptr);
                
                void getY(double *y_ptr)const;
                        
                void getYDot(double *ydot_ptr)const;
                
                /** Get number of Q's associated with body specified by BodyIndex. */
                const int getBodyNQ(const BodyIndex& B) const;
                
                /** Get number of U's associated with body specified by BodyIndex. */
                const int getBodyNU(const BodyIndex& B) const;
                
                /** Get number of Q's associated with body specified by BodyIndex. */
                const int getBodyNQ(const int& NodeIndex) const;
                
                /** Get number of U's associated with body specified by BodyIndex. */
                const int getBodyNU(const int& NodeIndex) const;
                
                /** const pointer to *Q of body B */
                const double *getQ() const;
                
                /** const pointer to *U of body B */
                const double *getU() const;                                
                
                const int getNY()const;
                
                double getTime()const;
                
                void addSysData(SystemData *ptr){SysData_ = ptr;}
                
                const SystemData& getSysData()const{return *SysData_;}
                
                /** Avoid using this function. This will be removed from future versions*/
                void setTime(const double& t){time_->setTime(t);}
                
                void initTime(TIME *time);
                
                void updQ(const double *q,const int& NodeIndex);
                
                void updQ(const double *q, const BodyIndex& BI);
                
                void updOneQ(const double& q, const int& NodeIndex, const QIndex& QI);
                
                void updU(const double *u,const int& NodeIndex);
                
                void updU(const double *u, const BodyIndex& BI);
                
                void updOneU(const double& u, const int& NodeIndex, const UIndex& UI);
                
                void updStateFromPrescribedMotion();
                
                // debug function
                void printQ();
                void printQDot();
                void printU();
                
                void printUDot();
                void displayUDot(){printUDot();}
                
                void printUDot1();
                
}; // class State

inline const int State::getNQ() const {return nQ_;}

inline const int State::getNU() const {return nU_;}

inline const double* State::getQ()const {return Q_;}

inline const double* State::getU()const {return U_;}

inline double* State::wgetQForNode(const int& NodeIndex){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Q_ + QIndexForNode_(NodeIndex);
        else
                return NULL;
                        
}

inline double* State::wgetQDotForNode(const int& NodeIndex){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return QDot_ + QIndexForNode_(NodeIndex);
        else
                return NULL;
                        
}

inline double* State::wgetUForNode(const int& NodeIndex){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return U_ + UIndexForNode_(NodeIndex);
        else return NULL;
}

inline double* State::wgetUDotForNode(const int& NodeIndex){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return UDot_ + UIndexForNode_(NodeIndex);
        else
                return NULL;
}

inline const double* State::getQForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Q_ + QIndexForNode_.getNode(NodeIndex);
        else
                return NULL;
}

//////////////////////////////////////////////////////////

inline Vector State::getQAsVectorForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Vector(BodyNQ_[NodeIndex],Q_ + QIndexForNode_.getNode(NodeIndex));
        else
                return Vector();
}

//////////////////////////////////////////////////////////

inline Vector State::getSystemQAsVector()const{
        if(nQ_)
                return Vector(nQ_,Q_);
        else
                return Vector();        
}

//////////////////////////////////////////////////////////

inline const Vector State::getQDotAsVectorForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Vector(BodyNQ_[NodeIndex],QDot_ + QIndexForNode_.getNode(NodeIndex));
        else
                return Vector();
}

//////////////////////////////////////////////////////////

inline const Vector State::getUAsVectorForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Vector(BodyNU_[NodeIndex],U_ + UIndexForNode_.getNode(NodeIndex));
        else
                return Vector();
}

//////////////////////////////////////////////////////////

inline const Vector State::getUDotAsVectorForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return Vector(BodyNU_[NodeIndex],UDot_ + UIndexForNode_.getNode(NodeIndex));
        else
                return Vector();
}

//////////////////////////////////////////////////////////

inline const double* State::getQDotForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return QDot_ + QIndexForNode_.getNode(NodeIndex);
        else
                return NULL;
}

//////////////////////////////////////////////////////////

inline const double* State::getUForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return U_ + UIndexForNode_.getNode(NodeIndex);
        else
                return NULL;
}

//////////////////////////////////////////////////////////

inline const double* State::getUDotForNode(const int& NodeIndex)const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        if(BodyNQ_[NodeIndex])
                return UDot_ + UIndexForNode_.getNode(NodeIndex);
        else
                return NULL;
}

//////////////////////////////////////////////////////////

inline const int State::getBodyNQ(const int& NodeIndex) const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        return BodyNQ_[NodeIndex];
}

//////////////////////////////////////////////////////////
                
inline const int State::getBodyNU(const int& NodeIndex) const{
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        return BodyNU_[NodeIndex];
}

//////////////////////////////////////////////////////////

inline const int State::getBodyNQ(const BodyIndex& BI) const{
        #ifdef RSIM_SAFE
                assert(BI.id()>=0);
                assert(BI.id()<nNodes_);
        #endif
        return BodyNQ_[BI.id()];
}

//////////////////////////////////////////////////////////
                
inline const int State::getBodyNU(const BodyIndex& BI) const{
        #ifdef RSIM_SAFE
                assert(BI.id()>=0);
                assert(BI.id()<nNodes_);
        #endif
        return BodyNU_[BI.id()];
}

//////////////////////////////////////////////////////////

inline void State::updOneQ(const double& q, const int& NodeIndex, const QIndex& QI){
        #ifdef RSIM_SAFE
                Utils::rsim_assert(NodeIndex>=0,
                                "State::updOneQ(), NodeIndex must be greater than zero.\n");
                                
                Utils::rsim_assert(NodeIndex<nNodes_,
                                "State::updOneQ(), NodeIndex too high.\n");
                Utils::rsim_assert(QI.id()>=0,
                                "State::updOneQ(), QIndex must be greater than zero.\n");
                Utils::rsim_assert(QI.id()<=BodyNQ_[NodeIndex],
                                "State::updOneQ(), QIndex too high.\n");
        #endif
        *(Q_+QIndexForNode_(NodeIndex)+QI.id()) = q;
}

//////////////////////////////////////////////////////////

inline void State::updOneU(const double& u, const int& NodeIndex, const UIndex& UI){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
                assert(UI.id()>=0);
                assert(UI.id()<=BodyNU_[NodeIndex]);
        #endif
        *(U_+UIndexForNode_(NodeIndex)+UI.id()) = u;
}

//////////////////////////////////////////////////////////

inline void State::updU(const double *u, const BodyIndex& BI){
        #ifdef RSIM_SAFE
                assert(BI.id()>=0);
                assert(BI.id()<nNodes_);
        #endif
        this->updU(u,BI.id());
}

//////////////////////////////////////////////////////////

// inline void State::updStateFromPrescribedMotion(){
//         for(int i=0;i<PrescribedMotionIndex_.getLength();++i){
//                 cout<<"\nHere";
//                 cin.get();
//                 int& index = PrescribedMotionIndex_[i];
//                 PM_(index)->setQ(Q_+QIndexForNode_(index));
//                 PM_(index)->setU(U_+UIndexForNode_(index));
//         }
// }

//////////////////////////////////////////////////////////

// inline double& State::Time(){
//         return time_->wgetTime();
// }

//////////////////////////////////////////////////////////

// inline const double& State::getTime()const{
//         return time_->getTime();
// }

//////////////////////////////////////////////////////////

inline const double* State::getQDot()const{
        return QDot_;
}

//////////////////////////////////////////////////////////
                
inline const double* State::getUDot()const{
        return UDot_;
}

//////////////////////////////////////////////////////////

inline void State::setY(const double *y_ptr){	
        for(int i=0;i<nQ_;++i)
                Q_[i]=y_ptr[i];
        
        for(int j=0;j<nU_;++j)
                U_[j]=y_ptr[nQ_+j];
}

//////////////////////////////////////////////////////////

inline void State::getYDot(double *ydot_ptr)const{
        for(int i=0;i<nQ_;++i)
                ydot_ptr[i] = QDot_[i];
        for(int j=0;j<nU_;++j)
                ydot_ptr[nQ_+j]=UDot_[j];
}

//////////////////////////////////////////////////////////

inline void State::getY(double *y_ptr)const{
        for(int i=0;i<nQ_;++i)
                y_ptr[i] = Q_[i];
        for(int j=0;j<nU_;++j)
                y_ptr[nQ_+j]=U_[j];
}

//////////////////////////////////////////////////////////

inline const int State::getNY()const{
        return (nQ_ + nU_);
}

//////////////////////////////////////////////////////////

} // namespace RSIM

#endif